import pandas as pd
import matplotlib.pyplot as plt
import re
import os


def main():
    if not os.path.exists('images'):
        os.makedirs('images')
    df = pd.read_csv("csv/dalian_weather_2022-2024_full.csv")

    # 数据清洗
    df['日期'] = pd.to_datetime(df['日期'])
    df['月份'] = df['日期'].dt.month
    df['年份'] = df['日期'].dt.year

    # 标准化天气描述
    def standardize_weather_type(weather_str):
        pattern = r'(小雨|中雨|大雨|暴雨|阵雨|雷阵雨|小到中雨|中到大雨|大到暴雨|阵雪|雨夹雪|小雪|中雪|大雪|暴雪|多云|阴|晴)'
        match = re.search(pattern, weather_str)

        if match:
            weather_key = match.group(1)
            if '雨' in weather_key or '雨夹雪' in weather_key:
                return '雨天'
            elif '雪' in weather_key:
                return '雪天'
            elif '晴' in weather_key:
                return '晴天'
            elif '阴' in weather_key:
                return '阴天'
            elif '多云' in weather_key:
                return '多云'
        print(weather_str)
        return '未知'

    df['白天天气状况'] = df['白天天气'].apply(standardize_weather_type)
    df['夜间天气状况'] = df['夜间天气'].apply(standardize_weather_type)

    def count_weather_avg(data, weather_column):
        yearly_counts = data.groupby(['年份', '月份', weather_column]).size().unstack().fillna(0)
        monthly_avg = yearly_counts.groupby('月份').mean()
        return monthly_avg

    day_weather_avg = count_weather_avg(df, '白天天气状况')
    night_weather_avg = count_weather_avg(df, '夜间天气状况')
    combined_weather_avg = day_weather_avg.add(night_weather_avg)

    # 设置中文显示
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False

    # 定义月份名称
    month_names = ['1月', '2月', '3月', '4月', '5月', '6月', '7月', '8月', '9月', '10月', '11月', '12月']

    # 获取所有天气状况并排序
    all_levels = sorted(set(day_weather_avg.columns) | set(night_weather_avg.columns),
                        key=lambda x: int(re.search(r'\d+', x).group(0)) if re.search(r'\d+', x) else 0)

    # 为所有天气状况创建颜色映射
    def create_weather_colormap(levels):
        cmap = plt.get_cmap('Pastel1')
        colors = {}
        level_count = len(levels)
        for i in range(level_count):
            ratio = float(i) / level_count
            color_value = cmap(ratio)
            current_level = levels[i]
            colors[current_level] = color_value
        return colors

    color_map = create_weather_colormap(all_levels)

    # 绘制柱状图
    def plot_weather_bar(data, title, filename):
        plt.figure(figsize=(16, 14))
        plt.suptitle(title, fontsize=22, y=0.98)

        # 获取所有天气类型
        weather_types = sorted(data.columns)

        for month in range(1, 13):
            plt.subplot(4, 3, month)
            if month in data.index:
                month_data = data.loc[month]
                month_data = month_data[month_data > 0]

                if month_data.empty:
                    plt.text(0.5, 0.5, f'{month_names[month - 1]}\n无数据',
                             ha='center', va='center', fontsize=12)
                    plt.axis('off')
                    continue

                # 为每种天气类型分配颜色
                colors = [color_map[wt] for wt in month_data.index]

                # 绘制柱状图
                bars = plt.bar(month_data.index, month_data.values, color=colors, edgecolor='grey')

                # 添加数值标签
                for bar in bars:
                    height = bar.get_height()
                    plt.annotate(f'{height:.1f}',
                                 xy=(bar.get_x() + bar.get_width() / 2, height),
                                 xytext=(0, 3),  # 垂直偏移
                                 textcoords="offset points",
                                 ha='center', va='bottom',
                                 fontsize=9)

                plt.title(f'{month_names[month - 1]}', fontsize=14, pad=10)
                plt.ylim(0, max(month_data.max() * 1.2, 5))

                # 旋转X轴标签
                plt.xticks(rotation=15, fontsize=9)
                plt.ylabel('平均天数', fontsize=10)
            else:
                plt.text(0.5, 0.5, f'{month_names[month - 1]}\n无数据',
                         ha='center', va='center', fontsize=12)
                plt.axis('off')

        # 添加图例
        handles = [plt.Rectangle((0, 0), 1, 1, color=color_map[wt]) for wt in weather_types]
        plt.figlegend(handles, weather_types,
                      loc='lower center',
                      ncol=min(6, len(weather_types)),
                      title='天气类型',
                      fontsize=11,
                      title_fontsize=12,
                      bbox_to_anchor=(0.5, 0.01))

        plt.tight_layout(rect=(0, 0.05, 1, 0.95))
        plt.savefig(f'../weather/images/{filename}', dpi=300, bbox_inches='tight')
        plt.show()

    # 使用示例
    plot_weather_bar(day_weather_avg,
                     '大连近三年各月白天天气分布 (2022-2024)',
                     'dalian_day_weather_bar.png')

    plot_weather_bar(night_weather_avg,
                     '大连近三年各月夜间天气分布 (2022-2024)',
                     'dalian_night_weather_bar.png')

    plot_weather_bar(combined_weather_avg,
                     '大连近三年各月综合天气分布 (2022-2024)',
                     'dalian_combined_weather_bar.png')

    print("柱状图已保存到../weather/images/目录")


if __name__ == '__main__':
    main()